#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 

class KernelAddcmul {
public:
    __aicore__ inline KernelAddcmul() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR x1, GM_ADDR x2, GM_ADDR value, GM_ADDR y, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t formerNum,
                              uint32_t formerLength, uint32_t formertileNum,
                              uint32_t formertileLength,
                              uint32_t formerlasttileLength, uint32_t tailNum,
                              uint32_t tailLength, uint32_t tailtileNum,
                              uint32_t tailtileLength,
                              uint32_t taillasttileLength, uint32_t tilingKey)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        if (tilingKey == 1) {
            this->blockLength = blockLength;
            this->tileNum =
                tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
            this->tileLength = tileLength / BUFFER_NUM;
            this->lasttileLength = lasttileLength;

            xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->blockLength * GetBlockIdx(),
                                this->blockLength);
            x1Gm.SetGlobalBuffer((__gm__ DTYPE_X*)x1 + this->blockLength * GetBlockIdx(),
                                this->blockLength);
            x2Gm.SetGlobalBuffer((__gm__ DTYPE_X*)x2 + this->blockLength * GetBlockIdx(),
                                this->blockLength);
            vGm.SetGlobalBuffer((__gm__ DTYPE_X*)value + this->blockLength * GetBlockIdx(),
                                this->blockLength);
            yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y + this->blockLength * GetBlockIdx(),
                                this->blockLength);
        }

        Init_v((DTYPE_X)0);

        if (tilingKey == 2) {
            this->formerNum = formerNum;
            this->formerLength = formerLength;
            this->formertileNum = formertileNum;
            this->formertileLength = formertileLength;
            this->formerlasttileLength = formerlasttileLength;

            this->tailNum = tailNum;
            this->tailLength = tailLength;
            this->tailtileNum = tailtileNum;
            this->tailtileLength = tailtileLength;
            this->taillasttileLength = taillasttileLength;

            if (GetBlockIdx() < this->formerNum) { 
                this->tileLength = this->formertileLength / BUFFER_NUM;
                this->lasttileLength = this->formerlasttileLength;
                this->tileNum = this->formertileNum * BUFFER_NUM;
                xGm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)x + this->formerLength * GetBlockIdx(),
                    this->formerLength);
                x1Gm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)x1 + this->formerLength * GetBlockIdx(),
                    this->formerLength);
                x2Gm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)x2 + this->formerLength * GetBlockIdx(),
                    this->formerLength);
                vGm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)value + this->formerLength * GetBlockIdx(),
                    this->formerLength);
                yGm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)y + this->formerLength * GetBlockIdx(),
                    this->formerLength);
            } else {  
                this->tileLength = this->tailtileLength / BUFFER_NUM;
                this->lasttileLength = this->taillasttileLength;
                this->tileNum = this->tailtileNum * BUFFER_NUM;
                xGm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)x + this->formerLength * this->formerNum +
                        this->tailLength * (GetBlockIdx() - this->formerNum),
                    this->tailLength);
                x1Gm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)x1 + this->formerLength * this->formerNum +
                        this->tailLength * (GetBlockIdx() - this->formerNum),
                    this->tailLength);
                x2Gm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)x2 + this->formerLength * this->formerNum +
                        this->tailLength * (GetBlockIdx() - this->formerNum),
                    this->tailLength);
                vGm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)value + this->formerLength * this->formerNum +
                        this->tailLength * (GetBlockIdx() - this->formerNum),
                    this->tailLength);
                yGm.SetGlobalBuffer(
                    (__gm__ DTYPE_X*)y + this->formerLength * this->formerNum +
                        this->tailLength * (GetBlockIdx() - this->formerNum),
                    this->tailLength);
            }
        }

        pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 3 * sizeof(DTYPE_X));
        pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
    }

    template <typename T>
    __aicore__ inline void Init_v(T flag) {
        this->value = vGm.GetValue(0);
    }
    template <>
    __aicore__ inline void Init_v(int8_t flag) {
        pipe.InitBuffer(tempBuf1, this->tileLength * sizeof(half));
        pipe.InitBuffer(tempBuf2, this->tileLength * sizeof(half));
        pipe.InitBuffer(tempBuf3, this->tileLength * sizeof(half));
        pipe.InitBuffer(tempBuf4, this->tileLength * sizeof(half));
        pipe.InitBuffer(tempBuf5, this->tileLength * sizeof(int16_t));
        pipe.InitBuffer(tempBuf6, this->tileLength * sizeof(int16_t));
        pipe.InitBuffer(tempBuf7, this->tileLength * sizeof(int16_t));
        pipe.InitBuffer(tempBuf8, this->tileLength * sizeof(int16_t));
    }
    template <>
    __aicore__ inline void Init_v(int32_t flag) {
        this->value_32 = vGm.GetValue(0);
    }

    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++) {
        CopyIn(i);
        Compute(i, this->value);
        CopyOut(i);
        }
    }

private:

    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_X> inLocal = inQueueIN.AllocTensor<DTYPE_X>();

        if (BUFFER_NUM == 1) {
        if (progress == this->tileNum - 1) {
            if (progress == 0) {
            DataCopy(inLocal[0], xGm[0], this->tileLength);
            DataCopy(inLocal[this->tileLength], x1Gm[0], this->tileLength);
            DataCopy(inLocal[2 * (this->tileLength)], x2Gm[0], this->tileLength);
            } else {
            DataCopy(
                inLocal[0],
                xGm[(progress - 1) * this->tileLength + this->lasttileLength],
                this->tileLength);
            DataCopy(
                inLocal[this->tileLength],
                x1Gm[(progress - 1) * this->tileLength + this->lasttileLength],
                this->tileLength);
            DataCopy(
                inLocal[2 * (this->tileLength)],
                x2Gm[(progress - 1) * this->tileLength + this->lasttileLength],
                this->tileLength);
            }
        } else {
            DataCopy(inLocal[0], xGm[progress * this->tileLength],
                    this->tileLength);
            DataCopy(inLocal[this->tileLength], x1Gm[progress * this->tileLength],
                    this->tileLength);
            DataCopy(inLocal[2 * (this->tileLength)],
                    x2Gm[progress * this->tileLength], this->tileLength);
        }
        }
        if (BUFFER_NUM == 2) {
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
            (progress == (this->tileNum * BUFFER_NUM - 1))) {
            DataCopy(
                inLocal[0],
                xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                (this->tileLength));
            DataCopy(
                inLocal[this->tileLength],
                x1Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                (this->tileLength));
            DataCopy(
                inLocal[2 * (this->tileLength)],
                x2Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                (this->tileLength));
        }

        else {
            DataCopy(inLocal[0], xGm[progress * (this->tileLength)],
                    (this->tileLength));
            DataCopy(inLocal[this->tileLength], x1Gm[progress * this->tileLength],
                    this->tileLength);
            DataCopy(inLocal[2 * (this->tileLength)],
                    x2Gm[progress * this->tileLength], this->tileLength);
        }
        }

        inQueueIN.EnQue(inLocal);
    }

    template <typename T>
    __aicore__ inline void Compute(int32_t progress, T flag)
    {
        LocalTensor<T> inLocal = inQueueIN.DeQue<T>();
        LocalTensor<T> xLocal = inLocal;
        LocalTensor<T> x1Local = inLocal[this->tileLength];
        LocalTensor<T> x2Local = inLocal[2 * (this->tileLength)];

        LocalTensor<T> yLocal = outQueueOUT.AllocTensor<T>();

        Mul(yLocal, x1Local, x2Local, this->tileLength);
        Muls(yLocal, yLocal, this->value, this->tileLength);
        Add(yLocal, yLocal, xLocal, this->tileLength);
        outQueueOUT.EnQue<T>(yLocal);
        inQueueIN.FreeTensor(inLocal);
    }

    template <>
    __aicore__ inline void Compute(int32_t progress, int8_t flag)
    {
        LocalTensor<int8_t> inLocal = inQueueIN.DeQue<int8_t>();
        LocalTensor<int8_t> xLocal = inLocal;
        LocalTensor<int8_t> x1Local = inLocal[this->tileLength];
        LocalTensor<int8_t> x2Local = inLocal[2 * (this->tileLength)];

        LocalTensor<half> temp1 = tempBuf1.Get<half>();
        LocalTensor<half> temp2 = tempBuf2.Get<half>();
        LocalTensor<half> temp3 = tempBuf3.Get<half>();
        LocalTensor<half> temp4 = tempBuf4.Get<half>();
        LocalTensor<int16_t> temp5 = tempBuf5.Get<int16_t>();
        LocalTensor<int16_t> temp6 = tempBuf1.Get<int16_t>();
        LocalTensor<int16_t> temp7 = tempBuf2.Get<int16_t>();
        LocalTensor<int16_t> temp8 = tempBuf3.Get<int16_t>();
        LocalTensor<int8_t> yLocal = outQueueOUT.AllocTensor<int8_t>();


        if (progress == 0) {
            Cast(temp4, vLocal, RoundMode::CAST_NONE, 1);
            Cast(temp8, temp4, RoundMode::CAST_RINT, 1);
            this->value_16 = temp8.GetValue(0);
        }


        Cast(temp1, xLocal, RoundMode::CAST_NONE, this->tileLength);
        Cast(temp2, x1Local, RoundMode::CAST_NONE, this->tileLength);
        Cast(temp3, x2Local, RoundMode::CAST_NONE, this->tileLength);

        Cast(temp5, temp1, RoundMode::CAST_RINT, this->tileLength);
        Cast(temp6, temp2, RoundMode::CAST_RINT, this->tileLength);
        Cast(temp7, temp3, RoundMode::CAST_RINT, this->tileLength);

        Mul(temp7, temp7, temp6, this->tileLength);
        Muls(temp7, temp7, this->value_16, this->tileLength);
        Add(temp7, temp7, temp5, this->tileLength);
        
        Cast(temp4, temp7, RoundMode::CAST_NONE, this->tileLength);
        Cast(yLocal, temp4, RoundMode::CAST_NONE, this->tileLength);

        outQueueOUT.EnQue<int8_t>(yLocal);
        inQueueIN.FreeTensor(inLocal);
    }

    template <>
    __aicore__ inline void Compute(int32_t progress, int32_t flag)
    {
        LocalTensor<int32_t> inLocal = inQueueIN.DeQue<int32_t>();
        LocalTensor<int32_t> xLocal = inLocal;
        LocalTensor<int32_t> x1Local = inLocal[this->tileLength];
        LocalTensor<int32_t> x2Local = inLocal[2 * (this->tileLength)];

        LocalTensor<int32_t> yLocal = outQueueOUT.AllocTensor<int32_t>();

        Mul(yLocal, x1Local, x2Local, this->tileLength);
        Muls(yLocal, yLocal, this->value_32, this->tileLength);
        Add(yLocal, yLocal, xLocal, this->tileLength);
        outQueueOUT.EnQue<int32_t>(yLocal);
        inQueueIN.FreeTensor(inLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
       LocalTensor<DTYPE_Y> outLocal = outQueueOUT.DeQue<DTYPE_Y>();

        if (BUFFER_NUM == 1) {
        if (progress == this->tileNum - 1) {
            if (progress == 0) {
            DataCopy(yGm[0], outLocal, this->tileLength);
            } else {
            DataCopy(
                yGm[(progress - 1) * this->tileLength + this->lasttileLength],
                outLocal, this->tileLength);
            }
        } else {
            DataCopy(yGm[progress * this->tileLength], outLocal,
                    this->tileLength);
        }
        }
        if (BUFFER_NUM == 2) {
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
            (progress == (this->tileNum * BUFFER_NUM - 1))) {
            DataCopy(
                yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                outLocal, (this->tileLength));
        }

        else {
            DataCopy(yGm[progress * (this->tileLength)], outLocal,
                    (this->tileLength));
        }
        }

        outQueueOUT.FreeTensor(outLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_X> vGm;
    GlobalTensor<DTYPE_Y> yGm;

    TBuf<> tempBuf1;
    TBuf<> tempBuf2;
    TBuf<> tempBuf3;
    TBuf<> tempBuf4;
    TBuf<> tempBuf5;
    TBuf<> tempBuf6;
    TBuf<> tempBuf7;
    TBuf<> tempBuf8;

    DTYPE_X value;
    int32_t value_32;
    int8_t  value_8;
    int16_t value_16;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;
    uint32_t formerNum;
    uint32_t formerLength;
    uint32_t formertileNum;
    uint32_t formertileLength;
    uint32_t formerlasttileLength;
    uint32_t tailNum;
    uint32_t tailLength;
    uint32_t tailtileNum;
    uint32_t tailtileLength;
    uint32_t taillasttileLength;
};

extern "C" __global__ __aicore__ void addcmul(GM_ADDR x, GM_ADDR x1, GM_ADDR x2, GM_ADDR value, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelAddcmul op;
    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        tilingKey = 2;
    } else {
        tilingKey = 1;
    }
    op.Init(x, x1, x2, value, y, 
          tiling_data.blockLength,
          tiling_data.tileNum, tiling_data.tileLength,
          tiling_data.lasttileLength, tiling_data.formerNum,
          tiling_data.formerLength, tiling_data.formertileNum,
          tiling_data.formertileLength, tiling_data.formerlasttileLength,
          tiling_data.tailNum, tiling_data.tailLength, tiling_data.tailtileNum,
          tiling_data.tailtileLength, tiling_data.taillasttileLength,
          tilingKey);
    op.Process();
}